import numpy as np
import cv2
import pickle
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

# 加载图像和标签
def load_image_and_label(image_path, label_path):
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
    return image, label

# 提取特征和标签
def create_features_and_labels(image, label):
    # 定义滤波器
    sobel_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)
    sobel_edges = np.sqrt(sobel_x**2 + sobel_y**2)

    # 使用Canny边缘检测
    canny_edges = cv2.Canny(image, 50, 150)

    # 将图像和边缘检测结果展平成一维数组
    X_image = image.reshape(-1, 1)
    X_sobel_edges = sobel_edges.reshape(-1, 1)
    X_canny_edges = canny_edges.reshape(-1, 1)

    # 合并新的特征列
    X_combined = np.concatenate((X_image, X_sobel_edges, X_canny_edges), axis=1)

    # 将标签展平并转换为二进制（0和1）
    y = label.reshape(-1)

    return X_combined, y

# 主函数
def main():
    image_path = 'C:/Users/86178/Desktop/Segment_sandstone/Sandstone_1.tif'
    label_path = 'C:/Users/86178/Desktop/Segment_sandstone/Sandstone_1_segment.tif'
    
    image, label = load_image_and_label(image_path, label_path)
    print(f"图像形状: {image.shape}")
    
    X, y = create_features_and_labels(image, label)
    
    print("完成从砂岩截面图1及其对应分区中获取X和y")
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
    
    print("完成train_test_split")
    
    # 训练随机森林分类器
    clf = RandomForestClassifier(n_estimators=100, random_state=42)
    clf.fit(X_train, y_train)
    
    # 在测试集上评估模型
    y_pred = clf.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"准确率: {accuracy}")
    
    # 保存模型到硬盘
    with open('clf.pkl', 'wb') as f:
        pickle.dump(clf, f)
    print("已保存随机森林模型clf到硬盘")

if __name__ == "__main__":
    main()